#!/usr/bin/env python3

import rclpy
from rclpy.node import Node
from rclpy.action import ActionClient
from nav_msgs.msg import Odometry
from std_msgs.msg import Float64
from nav2_msgs.action import NavigateToPose
from rclpy.executors import MultiThreadedExecutor
import time
import threading


class AutoFlightMission(Node):
    """自动飞行任务节点 - 使用简单z轴控制器"""
    
    def __init__(self):
        super().__init__('auto_flight_mission')
        
        # 创建导航action客户端
        self.nav_action_client = ActionClient(self, NavigateToPose, '/red_standard_robot1/navigate_to_pose')
        
        # 创建高度控制发布者
        self.altitude_pub = self.create_publisher(Float64, '/target_altitude', 10)
        
        # 订阅当前位置（用于监控）
        self.odom_sub = self.create_subscription(
            Odometry,
            '/red_standard_robot1/odometry',
            self.odometry_callback,
            10
        )
        
        # 状态变量
        self.current_odom = None
        self.mission_active = False
        
        self.get_logger().info('自动飞行任务节点已初始化（使用简单z轴控制器）')
        
        # 等待服务可用
        self.wait_for_services()
        
    def wait_for_services(self):
        """等待必要的服务可用"""
        self.get_logger().info('等待导航服务启动...')
        self.nav_action_client.wait_for_server()
        self.get_logger().info('导航服务已就绪')
        
        # 等待一段时间确保所有系统就绪
        time.sleep(2)
        
    def odometry_callback(self, msg):
        """更新当前位置"""
        self.current_odom = msg
        
    def set_altitude(self, target_altitude):
        """设置目标高度"""
        altitude_msg = Float64()
        altitude_msg.data = target_altitude
        
        # 发布几次确保收到
        for _ in range(3):
            self.altitude_pub.publish(altitude_msg)
            time.sleep(0.1)
            
        self.get_logger().info(f'设置目标高度: {target_altitude:.1f}m')
        
    def send_navigation_goal(self, x, y, yaw=0.0):
        """发送2D导航目标（只控制x,y）"""
        goal_msg = NavigateToPose.Goal()
        
        # 设置目标位置
        goal_msg.pose.header.frame_id = 'map'
        goal_msg.pose.header.stamp = self.get_clock().now().to_msg()
        goal_msg.pose.pose.position.x = x
        goal_msg.pose.pose.position.y = y
        goal_msg.pose.pose.position.z = 0.0  # z轴由独立控制器处理
        
        # 设置朝向（四元数）
        import math
        goal_msg.pose.pose.orientation.z = math.sin(yaw / 2.0)
        goal_msg.pose.pose.orientation.w = math.cos(yaw / 2.0)
        
        self.get_logger().info(f'发送2D导航目标: ({x:.1f}, {y:.1f})')
        
        # 发送目标
        future = self.nav_action_client.send_goal_async(goal_msg)
        return future
    
    def wait_for_navigation_complete(self, future, timeout=30.0):
        """等待导航完成"""
        start_time = time.time()
        
        # 等待目标被接受
        rclpy.spin_until_future_complete(self, future, timeout_sec=5.0)
        
        if not future.result():
            self.get_logger().error('导航目标被拒绝')
            return False
            
        goal_handle = future.result()
        if not goal_handle.accepted:
            self.get_logger().error('导航目标未被接受')
            return False
            
        self.get_logger().info('导航目标已接受，等待完成...')
        
        # 等待导航完成
        result_future = goal_handle.get_result_async()
        
        while not result_future.done():
            rclpy.spin_once(self, timeout_sec=0.1)
            
            # 检查超时
            if time.time() - start_time > timeout:
                self.get_logger().warn('导航超时')
                goal_handle.cancel_goal_async()
                return False
                
        result = result_future.result()
        if result.status == 4:  # SUCCEEDED
            self.get_logger().info('导航目标完成')
            return True
        else:
            self.get_logger().warn(f'导航失败，状态码: {result.status}')
            return False
    
    def wait_for_altitude_reached(self, target_altitude, tolerance=0.05, timeout=15.0):
        """等待高度到达目标"""
        start_time = time.time()
        stable_count = 0
        required_stable_count = 10  # 需要连续10次检查都在容差范围内
        
        self.get_logger().info(f'等待高度到达 {target_altitude:.1f}m (±{tolerance:.2f}m)...')
        
        while time.time() - start_time < timeout:
            if self.current_odom:
                current_alt = self.current_odom.pose.pose.position.z
                error = abs(current_alt - target_altitude)
                
                if error <= tolerance:
                    stable_count += 1
                    if stable_count >= required_stable_count:
                        self.get_logger().info(f'高度已到达: {current_alt:.3f}m')
                        return True
                else:
                    stable_count = 0
                    
                # 每2秒报告一次进度
                if int(time.time() - start_time) % 2 == 0:
                    self.get_logger().info(f'高度进度: 当前 {current_alt:.3f}m, 目标 {target_altitude:.1f}m, 误差 {error:.3f}m')
            else:
                if time.time() - start_time > 2.0:  # 2秒后还没有姿态数据就报警
                    self.get_logger().warn('未收到里程计数据，请检查/red_standard_robot1/odometry话题是否正常发布')
                    
            time.sleep(0.1)
            rclpy.spin_once(self, timeout_sec=0.01)
            
        self.get_logger().warn('高度到达超时')
        return False
    
    def wait_for_xy_position_reached(self, target_x, target_y, tolerance=0.2, timeout=30.0):
        """等待xy位置到达目标"""
        start_time = time.time()
        stable_count = 0
        required_stable_count = 5  # 需要连续5次检查都在容差范围内
        
        self.get_logger().info(f'等待到达xy位置: ({target_x:.1f}, {target_y:.1f}) (±{tolerance:.2f}m)')
        
        while time.time() - start_time < timeout:
            if self.current_odom:
                current_x = self.current_odom.pose.pose.position.x
                current_y = self.current_odom.pose.pose.position.y
                
                xy_error = ((current_x - target_x)**2 + (current_y - target_y)**2)**0.5
                
                if xy_error <= tolerance:
                    stable_count += 1
                    if stable_count >= required_stable_count:
                        self.get_logger().info(f'xy位置已到达: ({current_x:.2f}, {current_y:.2f}), 误差: {xy_error:.2f}m')
                        return True
                else:
                    stable_count = 0
                    # 每2秒报告一次进度
                    if int(time.time() - start_time) % 2 == 0:
                        self.get_logger().info(f'xy位置进度: 当前({current_x:.2f}, {current_y:.2f}), 误差: {xy_error:.2f}m')
            else:
                if time.time() - start_time > 2.0:
                    self.get_logger().warn('未收到里程计数据，请检查/red_standard_robot1/odometry话题是否正常发布')
                    
            time.sleep(0.1)
            rclpy.spin_once(self, timeout_sec=0.01)
            
        self.get_logger().warn('xy位置到达超时')
        return False
    
    def execute_mission(self):
        """执行飞行任务 - 简化版本"""
        if self.mission_active:
            self.get_logger().warn('任务已在运行中')
            return
            
        self.mission_active = True
        self.get_logger().info('=' * 50)
        self.get_logger().info('开始执行自动飞行任务（简单z轴控制版本）')
        self.get_logger().info('=' * 50)
        
        try:
            # 步骤1: 等待5秒并检查系统状态
            self.get_logger().info('步骤1: 等待5秒并检查系统状态...')
            time.sleep(3.0)
            
            # 检查里程计数据
            if not self.current_odom:
                self.get_logger().warn('未检测到里程计数据，请确保/red_standard_robot1/odometry话题正常发布')
                time.sleep(2.0)  # 再等待2秒
                
            time.sleep(2.0)
            
            # 步骤2: 起飞到0.8m高度
            self.get_logger().info('步骤2: 起飞到0.8m高度')
            self.set_altitude(0.8)
            if not self.wait_for_altitude_reached(0.8):
                raise Exception("起飞失败")
            time.sleep(1.0)  # 额外等待确保稳定
            
            # 步骤3: 飞到(2, 0)位置
            self.get_logger().info('步骤3: 飞到位置(2, 0)')
            future = self.send_navigation_goal(2.0, 0.0)
            if not self.wait_for_navigation_complete(future):
                raise Exception("导航到(2, 0)失败")
            # 等待xy位置到达
            if not self.wait_for_xy_position_reached(2.0, 0.0):
                self.get_logger().warn("xy位置可能未完全到达，但继续执行")
            time.sleep(1.0)
            
            # 步骤4: 降低到0.5m高度
            self.get_logger().info('步骤4: 降低到0.5m高度')
            self.set_altitude(0.5)
            if not self.wait_for_altitude_reached(0.5):
                raise Exception("降低高度失败")
            time.sleep(1.0)
            
            # 步骤5: 飞回(0, 0)位置
            self.get_logger().info('步骤5: 飞回初始位置(0, 0)')
            future = self.send_navigation_goal(0.0, 0.0)
            if not self.wait_for_navigation_complete(future):
                raise Exception("导航到(0, 0)失败")
            # 等待xy位置到达
            if not self.wait_for_xy_position_reached(0.0, 0.0):
                self.get_logger().warn("xy位置可能未完全到达，但继续执行")
            time.sleep(1.0)
            
            # 步骤6: 降低到0.2m高度
            self.get_logger().info('步骤6: 降低到0.2m高度')
            self.set_altitude(0.2)
            if not self.wait_for_altitude_reached(0.2):
                raise Exception("最终降低高度失败")
            
            self.get_logger().info('=' * 50)
            self.get_logger().info('🎉 自动飞行任务完成！')
            self.get_logger().info('=' * 50)
            
        except Exception as e:
            self.get_logger().error(f'任务执行失败: {str(e)}')
            self.get_logger().info('尝试紧急降落到0.2m...')
            self.set_altitude(0.2)
            
        finally:
            self.mission_active = False
    
    def start_mission(self):
        """在单独线程中启动任务"""
        if self.mission_active:
            self.get_logger().warn('任务已在运行中')
            return
            
        mission_thread = threading.Thread(target=self.execute_mission)
        mission_thread.daemon = True
        mission_thread.start()
        self.get_logger().info('任务线程已启动')


def main(args=None):
    rclpy.init(args=args)
    
    # 创建节点
    flight_mission = AutoFlightMission()
    
    # 使用多线程执行器
    executor = MultiThreadedExecutor()
    executor.add_node(flight_mission)
    
    # 启动任务
    flight_mission.start_mission()
    
    try:
        executor.spin()
    except KeyboardInterrupt:
        flight_mission.get_logger().info('收到退出信号')
    finally:
        flight_mission.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main() 